// Example 10-1. Pyramid Lucas-Kanade optical flow code
//
/* *************** License:**************************
 Oct. 3, 2008
 Right to use this code in any way you want without warrenty, support or any guarentee of it working.
 
 BOOK: It would be nice if you cited it:
 Learning OpenCV: Computer Vision with the OpenCV Library
 by Gary Bradski and Adrian Kaehler
 Published by O'Reilly Media, October 3, 2008
 
 AVAILABLE AT: 
 http://www.amazon.com/Learning-OpenCV-Computer-Vision-Library/dp/0596516134
 Or: http://oreilly.com/catalog/9780596516130/
 ISBN-10: 0596516134 or: ISBN-13: 978-0596516130    
 
 OTHER OPENCV SITES:
 * The source code is on sourceforge at:
 http://sourceforge.net/projects/opencvlibrary/
 * The OpenCV wiki page (As of Oct 1, 2008 this is down for changing over servers, but should come back):
 http://opencvlibrary.sourceforge.net/
 * An active user group is at:
 http://tech.groups.yahoo.com/group/OpenCV/
 * The minutes of weekly OpenCV development meetings are at:
 http://pr.willowgarage.com/wiki/OpenCV
 ************************************************** */

#include <cv.h>
#include <cxcore.h>
#include <highgui.h>
#include <stdio.h>

#include <iostream>

const int MAX_CORNERS = 10;
int main(int argc, char** argv) {

	// -- tracking set up-- //
	
	CvRect trackRect = cvRect( 100, 100, 100, 100 );
	
	int         win_size = 10;
	
	CvCapture* camera = cvCaptureFromCAM(-1);
	//cvCreateCameraCapture (CV_CAP_ANY);
	
	IplImage* currentFrame = cvQueryFrame(camera);
	
	CvSize      img_sz    = cvGetSize( currentFrame );
	
	// workspace
	IplImage* imgA = cvCreateImage( img_sz, IPL_DEPTH_8U, 1 );
	IplImage* imgB = cvCreateImage( img_sz, IPL_DEPTH_8U, 1 );
	
	cvConvertImage( currentFrame, imgB, CV_BGR2GRAY);
	//FIX: Confirm BGR is used initially
	
	// The first thing we need to do is get the features
	// we want to track.
	//
	IplImage* eig_image = cvCreateImage( img_sz, IPL_DEPTH_32F, 1 );
	IplImage* tmp_image = cvCreateImage( img_sz, IPL_DEPTH_32F, 1 );
	
	CvFont font;
	
	cvInitFont( &font, CV_FONT_HERSHEY_PLAIN, 1.0, 1.0);
	
	cvNamedWindow("Video",0);
	
	char* errString = new char[100]; //TODO correct size?
	
	
	// -- recording set up -- //
	
	//Could just use img_sz
	CvSize size = cvSize(
						 (int)cvGetCaptureProperty( camera, CV_CAP_PROP_FRAME_WIDTH),
						 (int)cvGetCaptureProperty( camera, CV_CAP_PROP_FRAME_HEIGHT)
						 );
    
	CvVideoWriter* writer = cvCreateVideoWriter(  // On linux Will only work if you've installed ffmpeg development files correctly, 
												argv[1],                               // otherwise segmentation fault.  Windows probably better.
												//CV_FOURCC('D', 'I', 'B', ' '), //BROKEN!
												CV_FOURCC('I', '4', '2', '0'),
												//0, // apparently works on Windows, but here? OSX : it seems to default to I420
												//CV_FOURCC('F', 'L', 'V', '1'),
												//CV_FOURCC('D','X','5','0'),    
												11, //fps unknown //TODO : what's the correct fps?
												size
												);

	IplImage* writeFrame = cvCreateImage(
										size,
										IPL_DEPTH_8U,
										3
										);
	
	// -- timing set up -- //
	
	int frameCount = 0;
	time_t	startTime = time (NULL);
	
	while ( (currentFrame = cvQueryFrame(camera)) ) {
		
		cvFlip( currentFrame, NULL, 1 );
		//NULL: Flip in place
		//1 : flip about y-axis
		
		IplImage* tmpPtr = imgA;
		imgA = imgB;
		imgB = tmpPtr;
		
		//TODO : change?
		IplImage* imgC = currentFrame;
		
		cvConvertImage( currentFrame, imgB, CV_BGR2GRAY);
		
		cvSetImageROI( imgA, cvRect( trackRect.x, trackRect.y, 60, 60 ) );
		
		
		int              corner_count = MAX_CORNERS;
		CvPoint2D32f* cornersA        = new CvPoint2D32f[ MAX_CORNERS ];
		cvGoodFeaturesToTrack(
							  imgA,
							  eig_image,
							  tmp_image,
							  cornersA,
							  &corner_count,
							  0.01,
							  10.0,
							  0,
							  3,
							  0,
							  0.04
							  );
		
		
		
		for( int i=0; i<corner_count; i++ ) {
			cornersA[i].x += trackRect.x;
			cornersA[i].y += trackRect.y;
		}
		
		cvResetImageROI( imgA ); //FIX: Inefficient! should not release image ROI
		
//		//JUST Checking location is correct
//		cvSetImageROI( imgC, cvRect( offsetX, offsetY, 30, 30 ) );
//		
//		cvSetZero( imgC );
//		
//		cvResetImageROI( imgC );
		
		
		// Call the Lucas Kanade algorithm
		//
		char features_found[ MAX_CORNERS ];
		float feature_errors[ MAX_CORNERS ];
		CvSize pyr_sz = cvSize( imgA->width+8, imgB->height/3 );
		IplImage* pyrA = cvCreateImage( pyr_sz, IPL_DEPTH_32F, 1 );
		IplImage* pyrB = cvCreateImage( pyr_sz, IPL_DEPTH_32F, 1 );
		CvPoint2D32f* cornersB        = new CvPoint2D32f[ MAX_CORNERS ];
		cvCalcOpticalFlowPyrLK(
							   imgA,
							   imgB,
							   pyrA,
							   pyrB,
							   cornersA,
							   cornersB,
							   corner_count,
							   cvSize( win_size,win_size ),
							   5,
							   features_found,
							   feature_errors,
							   cvTermCriteria( CV_TERMCRIT_ITER | CV_TERMCRIT_EPS, 20, .3 ),
							   0
							   );
		double sumDeltaX = 0;
		double sumDeltaY = 0;
		double invErrSum = 0;
		
		
		for( int i=0; i<corner_count; i++ ) {
			if( features_found[i]==0|| feature_errors[i]>550 ) {
				//       printf("Error is %f/n",feature_errors[i]);
				continue;
			}
			
			//    printf("Got it/n");
			CvPoint p0 = cvPoint(
								 cvRound( cornersA[i].x ),
								 cvRound( cornersA[i].y )
								 );
			CvPoint p1 = cvPoint(
								 cvRound( cornersB[i].x ),
								 cvRound( cornersB[i].y )
								 );
			
			double invErr = 1/(feature_errors[i]+1); //avoid division by zero (very unlikely anyway, but still)
			
			sumDeltaX += (p1.x - p0.x) * invErr;
			sumDeltaY += (p1.y - p0.y) * invErr;
			
			invErrSum += invErr;
			
			
			cvLine( imgC, p0, p1, CV_RGB(255,0,0),2 );
		
			std::string s;
			std::stringstream out;
			out << cvRound( feature_errors[i] );
			s = out.str();
			
			strcpy (errString, s.c_str());

			cvPutText( imgC, errString, p1, &font, cvScalar( 0, 255, 0 ) );
		}
		
		int weightedAvgDeltaX = sumDeltaX / invErrSum;
		int weightedAvgDeltaY = sumDeltaY / invErrSum;
		
		trackRect.x += weightedAvgDeltaX;
		trackRect.y += weightedAvgDeltaY;
		//FIXNOW: clip to inside frame
		
		if ( trackRect.x < 0 ) trackRect.x = 0;
		if ( imgC->width < trackRect.x + trackRect.width ) trackRect.x = imgC->width - trackRect.width;
		if ( trackRect.y < 0 ) trackRect.y = 0;
		if ( imgC->height < trackRect.y + trackRect.height ) trackRect.y = imgC->height - trackRect.height;
		
		cvRectangle( imgC, cvPoint(trackRect.x, trackRect.y), cvPoint(trackRect.x + trackRect.width, trackRect.y + trackRect.height), 
					cvScalarAll( 255 ) );
		
		cvCopy( imgC, writeFrame );
		
		cvWriteToAVI( writer, writeFrame );
		
		cvShowImage("Video",imgC);
		frameCount++;
		
		char c = cvWaitKey(1);
		if( c == 27 ) break;
	}
	
	time_t endTime = time (NULL);
	
	int elapsedTime = endTime - startTime;
	
	int avgFPS;
	
	if (elapsedTime != 0) {
		avgFPS = frameCount / elapsedTime;
	} else {
		avgFPS = -1;
	}
	
	delete errString;
	cvReleaseCapture( &camera ); 
	cvReleaseVideoWriter( &writer );
	
	printf("Avg fps =%d\n", avgFPS);

	return 0;
}

